package club.babyf.nativeai.service.impl;

import club.babyf.nativeai.config.NativeAiProperties;
import club.babyf.nativeai.entity.CustomHashMap;
import club.babyf.nativeai.entity.ToolCallAccumulator;
import club.babyf.nativeai.entity.chat.ChatChoice;
import club.babyf.nativeai.entity.chat.ChatCompletionResponse;
import club.babyf.nativeai.entity.chat.Message;
import club.babyf.nativeai.entity.chat.ToolCall;
import club.babyf.nativeai.functioncalling.invoker.FunctionClient;
import club.babyf.nativeai.functioncalling.registry.FunctionRegistry;
import club.babyf.nativeai.listener.CustomEventSourceListener;
import club.babyf.nativeai.service.IChatService;
import club.babyf.nativeai.utils.AsyncBaseQueue;
import club.babyf.nativeai.utils.JsonCustomUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.ContentType;
import cn.hutool.http.HttpUtil;
import com.fasterxml.jackson.core.type.TypeReference;
import com.google.common.net.HttpHeaders;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;
import org.apache.commons.lang.exception.ExceptionUtils;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * ChatServiceImpl.
 *
 * @author zhanglei.
 * @date 2025/2/18 10:43.
 * @description 问答服务实现.
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class ChatServiceImpl implements IChatService {

    private final NativeAiProperties nativeAiProperties;

    private final WebClient webClient;

    private final OkHttpClient okHttpClient;

    private final FunctionRegistry functionRegistry;

    private final FunctionClient functionClient;

    @Override
    public Object chat(String authorization, Boolean isCusFunCall, String body) {
        // 判断body是否为空
        if (StrUtil.isBlank(body)) {
            return "invalid body";
        }
        // 传入的apiKey
        String apiKey = authorization.split(" ")[1];
        // 配置的apiKey
        String apiKeyWithConfig = nativeAiProperties.getOpenai().getApiKey();
        // 最终使用的apiKey
        String finalApiKey = StrUtil.isBlank(apiKey) ? apiKeyWithConfig : apiKey;
        // 判断最终使用的apiKey是否为空
        if (StrUtil.isBlank(finalApiKey)) {
            return "invalid apiKey";
        }
        // baseUrl
        String baseUrl = nativeAiProperties.getOpenai().getBaseUrl();
        // 最终使用的url
        String finalUrl = baseUrl + "/v1/chat/completions";
        // 请求参数
        CustomHashMap<String, Object> bodyMap = JsonCustomUtil.parse(body, CustomHashMap.class);
        // 判断是否开启自定义函数调用
        if (!isCustomFunctionCalling(isCusFunCall)) {
            // 注册的函数参数
            registryFunctionParams(bodyMap);
        }
        // 是否流式输出
        Boolean stream = bodyMap.getBoolean("stream");
        // 构建请求参数
        String finalBody = JsonCustomUtil.toJsonString(bodyMap);
        log.info("请求参数：{}", finalBody);
        // 判断是否流式输出
        if (stream == null || !stream) {
            return sendReq(finalUrl, finalApiKey, isCusFunCall, finalBody);
        }
        return sendStreamReq(finalUrl, finalApiKey, isCusFunCall, finalBody);
    }


    /**
     * 非流式请求
     *
     * @param url          请求路径
     * @param apiKey       秘钥key
     * @param isCusFunCall 是否开启自定义函数调用
     * @param body         请求参数
     * @return 响应结果
     */
    private String sendReq(String url, String apiKey, Boolean isCusFunCall, String body) {
        // 发送请求
        String resBody = doSendReq(url, apiKey, body);
        // 判断是否开启自定义函数调用
        if (isCustomFunctionCalling(isCusFunCall)) {
            // 注册的函数参数
            return resBody;
        }
        // 判断是否需要调用函数
        boolean invokeFunction = functionClient.isInvokeFunction(resBody);
        if (!invokeFunction) {
            return resBody;
        }
        // 调用函数
        CustomHashMap<String, Object> reqBodyWithFunction = functionClient.invokeFunction(body, resBody);
        // 发送请求
        resBody = doSendReq(url, apiKey, JsonCustomUtil.toJsonString(reqBodyWithFunction));
        // 返回执行结果
        return resBody;
    }

    /**
     * 非流式请求
     *
     * @param url    请求路径
     * @param apiKey 秘钥key
     * @param body   请求参数
     * @return 响应结果
     */
    private String doSendReq(String url, String apiKey, String body) {
        // 发送请求
        return HttpUtil.createPost(url)
                .header("Content-Type", "application/json")
                .header("Authorization", "Bearer " + apiKey)
                .body(body)
                .execute()
                .body();
    }


    /**
     * 流式请求
     *
     * @param url          请求路径
     * @param apiKey       秘钥key
     * @param isCusFunCall 是否开启自定义函数调用
     * @param body         请求参数
     * @return 响应结果
     */
    public Flux<ServerSentEvent<String>> sendStreamReq(String url, String apiKey, Boolean isCusFunCall, String body) {
        // 获取数据流并订阅处理
        Flux<String> chatResponseStream = doSendStreamReq(url, apiKey, isCusFunCall, body);
        // 返回响应结果
        return chatResponseStream.map(chunk -> ServerSentEvent.builder(chunk).build());
    }

    /**
     * 流式请求
     *
     * @param url          请求路径
     * @param apiKey       秘钥key
     * @param isCusFunCall 是否开启自定义函数调用
     * @param body         请求参数
     * @return 响应结果
     */
    public Flux<String> doSendStreamReq(String url, String apiKey, Boolean isCusFunCall, String body) {
        // 获取数据流并订阅处理
        return getChatResponseStream(url, apiKey, isCusFunCall, body);
    }


    /**
     * 流式请求
     *
     * @param url    请求路径
     * @param apiKey 秘钥key
     * @param body   请求参数
     * @return 响应结果
     */
    private SseEmitter sendSSEStreamReq(String url, String apiKey, String body) {
        // 默认30秒超时, 设置为0L则永不超时
        SseEmitter sseEmitter = new SseEmitter(0L);
        // 完成后回调
        sseEmitter.onCompletion(() -> log.info("结束连接..................."));
        // 超时回调
        sseEmitter.onTimeout(() -> log.info("连接超时..................."));
        // 异常回调
        sseEmitter.onError(
                throwable -> {
                    try {
                        log.info("异常回调，异常原因：{}", ExceptionUtils.getFullStackTrace(throwable));
                        sseEmitter.send(SseEmitter.event()
                                .id(UUID.randomUUID().toString())
                                .name("发生异常")
                                .data(Message.builder().content("发生异常请重试").build())
                                .reconnectTime(3000));
                    } catch (IOException e) {
                        log.error("创建sse连接失败IO异常，异常原因：{}", ExceptionUtils.getFullStackTrace(e));
                        throw new RuntimeException("创建sse连接失败IO异常，异常原因：" + ExceptionUtils.getFullStackTrace(e));
                    }
                }
        );
        try {
            sseEmitter.send(SseEmitter.event().reconnectTime(5000));
        } catch (IOException e) {
            log.error("创建sse连接IO异常，失败原因：{}", ExceptionUtils.getFullStackTrace(e));
            throw new RuntimeException("创建sse连接IO异常，失败原因：" + ExceptionUtils.getFullStackTrace(e));
        }
        // 异步执行openai问答接口
        CompletableFuture.supplyAsync(() -> {
            doSendSSEStreamReq(url, apiKey, body, new CustomEventSourceListener(sseEmitter));
            return "execution completed";
        }, AsyncBaseQueue.SENDER_ASYNC).exceptionally(e -> {
            CompletionException completionException = new CompletionException(e);
            log.error("异步发起openai流式请求异常，异常原因：{}", ExceptionUtils.getFullStackTrace(completionException));
            throw completionException;
        });
        // 返回响应结果
        return sseEmitter;
    }

    /**
     * 执行流式请求
     *
     * @param url                 请求路径
     * @param apiKey              秘钥key
     * @param body                请求参数
     * @param eventSourceListener 事件源监听器
     */
    public void doSendSSEStreamReq(String url, String apiKey, String body, EventSourceListener eventSourceListener) {
        try {
            EventSource.Factory factory = EventSources.createFactory(this.okHttpClient);
            Request request = new Request.Builder()
                    .header("Authorization", "Bearer " + apiKey)
                    .url(url)
                    .post(RequestBody.Companion.create(body.getBytes(), MediaType.parse(ContentType.JSON.getValue())))
                    .build();
            // 创建事件
            factory.newEventSource(request, eventSourceListener);
        } catch (Exception e) {
            log.error("请求流式接口异常，异常原因：{}", ExceptionUtils.getFullStackTrace(e));
            throw new RuntimeException("请求流式接口异常，异常原因：" + ExceptionUtils.getFullStackTrace(e));
        }
    }


    /**
     * 注册函数参数
     *
     * @param bodyMap 请求参数
     */
    private void registryFunctionParams(CustomHashMap<String, Object> bodyMap) {
        List<CustomHashMap<String, Object>> functionParamList = functionRegistry.getFunctionParamList();
        if (CollUtil.isNotEmpty(functionParamList)) {
            bodyMap.put("tools", functionParamList);
        }
    }

    /**
     * 获取对话响应流
     *
     * @param url          url
     * @param apiKey       apiKey
     * @param isCusFunCall 是否开启自定义函数调用
     * @param body         body
     * @return 响应流
     */
    public Flux<String> getChatResponseStream(String url, String apiKey, Boolean isCusFunCall, String body) {
        StringBuilder contentBuffer = new StringBuilder();
        Map<Integer, ToolCallAccumulator> accumulators = new ConcurrentHashMap<>();
        return webClient.post()
                .uri(url)
                .header(HttpHeaders.CONTENT_TYPE, org.springframework.http.MediaType.APPLICATION_JSON_VALUE)
                .header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey)
                .bodyValue(body)
                .accept(org.springframework.http.MediaType.TEXT_EVENT_STREAM)
                .retrieve()
                .bodyToFlux(String.class)
                // 遇到结束标记停止
                .takeUntil("[DONE]"::equals)
                .flatMap(chunk -> processChunk(chunk, url, apiKey, isCusFunCall, body, contentBuffer, accumulators))
                .filter(info -> info != null && !StrUtil.isBlank(info));
    }

    /**
     * 处理数据块
     *
     * @param chunk         数据块
     * @param url           url
     * @param apiKey        apiKey
     * @param isCusFunCall  是否开启自定义函数调用
     * @param body          body
     * @param contentBuffer 内容缓冲区
     * @param accumulators  累积器
     * @return 响应流
     */
    private Flux<String> processChunk(String chunk, String url, String apiKey, Boolean isCusFunCall, String body, StringBuilder contentBuffer, Map<Integer, ToolCallAccumulator> accumulators) {
        if ("[DONE]".equals(chunk)) {
            return Flux.empty();
        }
        // 判断是否开启自定义函数调用
        if (isCustomFunctionCalling(isCusFunCall)) {
            // 如果开启自定义函数调用，直接返回chunk
            return Flux.just(chunk);
        }
        // 解析数据块为ChatCompletionResponse对象
        ChatCompletionResponse chunkObj = JsonCustomUtil.parse(chunk, ChatCompletionResponse.class);
        StringBuilder contentOfChunkBuffer = new StringBuilder();
        // 处理每个choice的内容和工具调用
        for (ChatChoice choice : chunkObj.getChoices()) {
            Message delta = choice.getDelta();
            // 收集文本内容
            if (delta.getContent() != null) {
                contentBuffer.append(delta.getContent());
                contentOfChunkBuffer.append(chunk);
            }
            if (choice.getFinishReason() != null) {
                contentOfChunkBuffer.append(chunk);
            }
            // 处理工具调用
            if (delta.getToolCalls() != null) {
                for (ToolCall toolCall : delta.getToolCalls()) {
                    int index = toolCall.getIndex();
                    ToolCallAccumulator acc = accumulators.computeIfAbsent(index, k -> new ToolCallAccumulator());
                    // 更新累积器
                    if (toolCall.getId() != null) acc.setId(toolCall.getId());
                    if (toolCall.getFunction() != null) {
                        if (toolCall.getFunction().getName() != null) {
                            acc.setName(toolCall.getFunction().getName());
                        }
                        if (toolCall.getFunction().getArguments() != null) {
                            acc.appendArguments(toolCall.getFunction().getArguments());
                        }
                    }
                }
            } else {
                if (CollUtil.isNotEmpty(accumulators)) {
                    // 获取累积器的索引集合
                    Set<Integer> indexSet = accumulators.keySet();
                    // 累积器的索引集合从小到大排序
                    List<Integer> indexList = new ArrayList<>(indexSet);
                    indexList.sort(Comparator.naturalOrder());
                    // 声明累积器列表
                    List<ToolCallAccumulator> accList = new ArrayList<>();
                    // 遍历累积器的索引集合
                    for (Integer index : indexList) {
                        ToolCallAccumulator acc = accumulators.get(index);
                        if (acc.isComplete()) {
                            accList.add(acc);
                        }
                    }
                    if (CollUtil.isNotEmpty(accList)) {
                        // 发射已收集的内容
                        return Flux.just(contentBuffer.toString()).concatWith(handleFunctionCall(url, apiKey, accList, isCusFunCall, body));
                    }
                }
            }
        }
        return Flux.just(contentOfChunkBuffer.toString());
    }

    /**
     * 处理函数调用
     *
     * @param url          url
     * @param apiKey       apiKey
     * @param accList      累积器列表
     * @param isCusFunCall 是否开启自定义函数调用
     * @param requestBody  请求体
     * @return 响应流
     */
    private Flux<String> handleFunctionCall(String url, String apiKey, List<ToolCallAccumulator> accList, Boolean isCusFunCall, String requestBody) {
        // 解析请求体
        CustomHashMap<String, Object> requestBodyMap = JsonCustomUtil.parse(requestBody, CustomHashMap.class);
        // 获取messages
        List<CustomHashMap<String, Object>> messages = requestBodyMap.getList("messages", new TypeReference<>() {
        });
        // 构建新的消息列表
        List<CustomHashMap<String, Object>> newMessages = new ArrayList<>(messages);
        // 遍历累积器列表
        for (ToolCallAccumulator acc : accList) {
            // 执行本地函数
            Object result = functionClient.doInvokeFunction(acc.getName(), acc.getArgumentsJson());
            // 构建assistantMessage
            CustomHashMap<String, Object> assistantMessage = new CustomHashMap<>();
            assistantMessage.put("role", "assistant");
            List<CustomHashMap<String, Object>> toolCallList = new ArrayList<>();
            CustomHashMap<String, Object> toolCall = new CustomHashMap<>();
            toolCall.put("id", acc.getId());
            toolCall.put("type", "function");
            CustomHashMap<String, Object> function = new CustomHashMap<>();
            function.put("name", acc.getName());
            function.put("arguments", acc.getArgumentsJson());
            toolCall.put("function", function);
            toolCallList.add(toolCall);
            assistantMessage.put("tool_calls", toolCallList);
            newMessages.add(assistantMessage);
            // 构建toolMessage
            CustomHashMap<String, Object> toolMessage = new CustomHashMap<>();
            toolMessage.put("role", "tool");
            toolMessage.put("tool_call_id", acc.getId());
            toolMessage.put("name", acc.getName());
            toolMessage.put("content", result);
            newMessages.add(toolMessage);
        }
        // 更新请求体
        requestBodyMap.put("messages", newMessages);
        // 递归调用获取新响应流
        return getChatResponseStream(url, apiKey, isCusFunCall, JsonCustomUtil.toJsonString(requestBodyMap));
    }


    /**
     * 是否开启自定义函数调用
     *
     * @param isCusFunCall 是否开启自定义函数调用
     * @return true/false
     */
    private boolean isCustomFunctionCalling(Boolean isCusFunCall) {
        return isCusFunCall != null && isCusFunCall;
    }

}
